from hat.callbacks import CallbackMixin
from hat.registry import OBJECT_REGISTRY

import nvtx

@OBJECT_REGISTRY.register
class NVTXCallback(CallbackMixin):
    def __init__(self, **kwargs):
        self.loop_rng = None
        self.epoch_rng = None
        self.step_rng = None
        self.batch_rng = None
        self.backward_rng = None
        self.optimizer_step_rng = None
        self.forward_rng = None
        self.dataloader_rng = None

    def on_loop_begin(self, **kwargs):
        msg = "train_loop_e{}_s{}".format(kwargs['num_epochs'], kwargs['num_steps'])
        self.loop_rng = nvtx.start_range(message=msg, color="green")
    
    def on_loop_end(self, **kwargs):
        if self.loop_rng is not None:
            nvtx.end_range(self.loop_rng)
    
    def on_epoch_begin(self, **kwargs):
        msg = "epoch_{}_train".format(kwargs['epoch_id'])
        self.epoch_rng = nvtx.start_range(message=msg, color="blue")

    def on_epoch_end(self, **kwargs):
        if self.epoch_rng is not None:
            nvtx.end_range(self.epoch_rng)
    
    def on_step_begin(self, **kwargs):
        msg = "step_{}_train".format(kwargs['step_id'])
        self.step_rng = nvtx.start_range(message=msg, color="yellow")

        mseg = 'datraloader_{}'.format(kwargs['step_id'])
        self.dataloader_rng = nvtx.start_range(message=mseg, color="blue")

    def on_step_end(self, **kwargs):
        if self.step_rng is not None:
            nvtx.end_range(self.step_rng)
    
    def on_batch_begin(self, **kwargs):
        if self.dataloader_rng is not None:
            nvtx.end_range(self.dataloader_rng)
        msg = "batch_train"
        self.batch_rng = nvtx.start_range(message=msg, color="purple")

    def on_batch_end(self, **kwargs):
        if self.batch_rng is not None:
            nvtx.end_range(self.batch_rng)

    def on_backward_begin(self, **kwargs):
        msg = "model_backward"
        self.backward_rng = nvtx.start_range(message=msg, color="cyan")

    def on_backward_end(self, **kwargs):
        if self.backward_rng is not None:
            nvtx.end_range(self.backward_rng)

    def on_optimizer_step_begin(self, **kwargs):
        msg = "optimizer_step"
        self.optimizer_step_rng = nvtx.start_range(message=msg, color="orange")

    def on_optimizer_step_end(self, **kwargs):
        if self.optimizer_step_rng is not None:
            nvtx.end_range(self.optimizer_step_rng)

    def on_forward_begin(self, **kwargs):
        msg = "model_forward"
        self.forward_rng = nvtx.start_range(message=msg, color="rapids")

    def on_forward_end(self, **kwargs):
        if self.forward_rng is not None:
            nvtx.end_range(self.forward_rng)
